package me.zhengjie.aspect;

import com.google.common.util.concurrent.RateLimiter;
import me.zhengjie.annotation.RateLimit;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.stereotype.Component;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.TimeUnit;

@Aspect
@Component
public class RateLimiterAspect {
    private final ConcurrentMap<String, RateLimiter> limiterMap = new ConcurrentHashMap<>();

    @Around("@annotation(rateLimit)")
    public Object around(ProceedingJoinPoint joinPoint, RateLimit rateLimit) throws Throwable {
        String methodName = joinPoint.getSignature().toString();
        RateLimiter limiter = limiterMap.computeIfAbsent(methodName,
                k -> RateLimiter.create(rateLimit.value()));

        if (rateLimit.timeout() > 0) {
            if (!limiter.tryAcquire(rateLimit.timeout(), TimeUnit.MILLISECONDS)) {
                throw new RuntimeException("请求过于频繁，请稍后在试！");
            }
        } else {
            limiter.acquire();
        }
        return joinPoint.proceed();
    }
}

